Speculative 路径发送 sampling_mask 时用的是 model_output.accept_num 来做分组并按 i 构造 mask_dict key，但上面 speculate_save_output(_topk) 的输出会经过 index_to_batch_id + enable_pd_reorder 恢复到原始 batch 顺序；如果开启 PD reorder，这里未对 sampler_output.sampling_mask / accept_num / logz_per_batch 做一致的恢复排序，mask_dict 的 key/分组将与 token_processor 侧的 batch_id 不一致。建议：在生成 mask_dict 前先对 accept_num 与 sampling_mask 做与输出一致的 recover/reorder（可复用 recover_share_inputs["accept_num_cpu"] 或扩展 recover_batch_index_for_sampler_output），并同步重排 logz_per_batch。

logprobs_renormalize_with_logz 目前对所有 isfinite 的位置统一做 logprobs - logZ_K，但 logprobs_tensors 里的 top-k 项是从“全量分布”topk 取出的，未必全部落在 top_p/top_k 截断后的候选集合 K 内（尤其当 top_p 很小且 max_logprobs 较大时）。这会导致返回的“重归一化 logprobs”仍包含候选集之外 token 的有限值，不符合截断分布语义。建议结合 sampling_mask（或 candidate set）把不在 K 内的 token logprobs 置为 -inf / None，并仅对 K 内条目做重归一化，或改为直接在截断后的分布上构造 logprobs 输出。

这里用 block=True 同步等待 sampling_mask side-channel 消息，缺少超时/降级路径：一旦 worker 未发送（例如 client 未创建/发送失败/某些 runner 未接入该 side-channel），TokenProcessor 会永久阻塞，导致整体推理挂死。建议改为非阻塞轮询（block=False）并在缺失时允许该 step 继续，或增加可配置超时并打印错误日志，避免死锁。

这里的 cfg/model_config 使用 MagicMock 时，TokenProcessor.init 里 getattr(cfg.model_config, "enable_keep_sampling_mask", False) 会返回一个 truthy 的 MagicMock，导致单测意外开启 keep_sampling_mask 并尝试 bind 固定的 IPC 地址（/dev/shm/sampling_mask_output_rank_0_9700.socket），容易在并行/重复运行时出现“Address already in use”或资源泄漏。建议在 mock_cfg.model_config 上显式设置 enable_keep_sampling_mask=False（或 patch envs.FD_USE_GET_SAVE_OUTPUT_V1=True 以避免创建该 server）。

该 setUp 使用 MagicMock 构造 cfg 时同样存在 enable_keep_sampling_mask 被 MagicMock 误判为 True 的风险，TokenProcessor 可能在单测中意外创建并 bind sampling_mask 的 ZMQ IPC socket，造成端口/文件冲突和测试不稳定。建议显式设置 self.mock_cfg.model_config.enable_keep_sampling_mask = False。

该测试 cfg 通过 MagicMock 构造，TokenProcessor 初始化时可能将 enable_keep_sampling_mask 读取为 truthy 的 MagicMock，从而在单测里意外创建并 bind sampling_mask 的 ZMQ IPC server（固定 name/端口），导致用例间冲突或资源泄漏。建议在 cfg.model_config 上显式设置 enable_keep_sampling_mask=False。

这里对 logprobs 做 renormalize 时需要避免与 Sampler.compute_logprobs 中的 top_p_normalized_logprobs 逻辑重复归一化；否则当请求侧已开启 top_p_normalized_logprobs（top_p!=1.0）时会出现二次减去 logZ，导致返回的 logprobs 数值错误。建议按 request/token 维度判断是否已做过 top_p 归一化，再决定是否应用 logz_per_batch（或仅对未归一化的行应用）。

_compute_sampling_mask() 里 boundary_idx 由 bool sum 得到的是 int32，直接传给 paddle.take_along_axis 可能触发索引 dtype 不兼容（Paddle 通常要求 int64 索引），导致启用 keep_sampling_mask 时运行时报错。建议在 take_along_axis 前显式把 boundary_idx cast 到 int64。

save_output_normal() 在 enable_pd_reorder=True 时会调用 recover_batch_index_for_sampler_output()，但该函数当前不会重排 sampler_output.sampling_mask（以及 logz_per_batch）。这样会导致 sampling_mask 与 recover 后的 sampled_token_ids / batch_id 对不上，返回给客户端的 sampling_mask 可能错配到其他 request。建议在 recover 流程里把 sampling_mask/logz_per_batch 也按 index_to_batch_id 同步重排，或在发送 mask_dict 前基于 index_to_batch_id 做一次 list 重排。

save_output_specualate() 里 sampling_mask 的分组/发送同样没有考虑 enable_pd_reorder 的 index_to_batch_id 重排（recover_batch_index_for_sampler_output 也不会处理 sampling_mask）。在开启 PD reorder 时这里会出现 per-request 的 sampling_mask 分发错位。建议在构造 mask_dict 前先对 sampling_mask 与 accept_num 的对齐关系做恢复（要么扩展 recover_batch_index_for_sampler_output 支持 sampling_mask/logz_per_batch，要么在这里显式按 index_to_batch_id 重排后再按 accept_num 分组）。

这里在发送 sampling_mask 时未检查 sampling_mask_zmq_client 是否为 None。当前参数声明允许为 None，一旦上游配置/注入不一致（例如 keep_sampling_mask 被置真但 client 未初始化）会直接 AttributeError 并中断推理线程。建议在 send_pyobj 前加显式判空/断言，并给出更清晰的错误信息。

SamplerOutput.logz_per_batch 的注释写的是“Shape: [num_reqs]”，但在 speculative decoding 路径里 logz_per_batch 实际是按 accepted token 展平计算的（shape 更接近 [total_accepted_tokens]），仅用于 logprobs 重归一化。建议更新注释/命名以反映两种路径的真实维度，避免后续误用。

StreamTransferData.sampling_mask 的类型注解写成了 Optional[np.array]，但 np.array 是函数而不是类型；这里应使用 np.ndarray（或更具体的 np.ndarray[np.int32] 等）。否则会误导静态检查/IDE。

🔴 Bug sampling_mask_zmq_client 参数类型为 Optional[ZmqIpcClient]（默认 None），但此处未做空值检查直接调用 send_pyobj。

当前逻辑仅检查 sampler_output.sampling_mask is not None and model_output.mp_rank == 0，如果配置不一致（例如 enable_keep_sampling_mask=True 但 ZMQ client 因某些原因未初始化），将触发 AttributeError: 'NoneType' object has no attribute 'send_pyobj'。

建议修复：

if sampler_output.sampling_mask is not None and model_output.mp_rank == 0 and sampling_mask_zmq_client is not None:

🔴 Bug 与 save_output_normal 中相同的问题：sampling_mask_zmq_client 可能为 None（参数默认值为 None），调用前需添加空值防护。

建议修复：

if sampler_output.sampling_mask is not None and model_output.mp_rank == 0 and sampling_mask_zmq_client is not None:

🟡 建议 当同时启用 --enable-logprob 和 --enable-keep-sampling-mask 时，此处会将 logprobs 隐式地基于候选集重归一化。这改变了现有 logprobs 的语义（从全词表 log_softmax 变为截断候选集归一化），可能影响已依赖原始 logprobs 值的下游用户。

建议：

1. 将重归一化作为独立的可选行为（例如 --logprobs-renormalize 参数），或
2. 在 API 响应中新增独立字段（如 normalized_logprobs）而非覆盖原 logprobs，或
3. 至少在文档/启动日志中明确说明 --enable-keep-sampling-mask 会修改 logprobs 输出语义。

🟡 建议 _compute_sampling_mask 在每个 decode step 对 [B, vocab_size] 张量执行 argsort（O(B·V·logV)），对于大词表模型（如 GLM-4 的 151k vocab）开销显著。

而 top_k_top_p_sampling 内部已经执行了类似的排序逻辑。考虑：

1. 复用 top_k_top_p_sampling 中已有的排序结果来构建 mask，避免重复排序；
2. 或将 mask 计算与采样合并为一个 kernel 调用。

此外，tuple[List[np.ndarray], np.ndarray] 使用了 Python 3.9+ 的小写 tuple 语法，如项目需兼容 3.8 应改为 Tuple[List[np.ndarray], np.ndarray]。